"""训练集过采样
建议只对2020数据集过采样，2019数据集过多不必再采样
"""
from collections import defaultdict
import os
import random
from tqdm import tqdm
Train2019_OverSample = False
Train2020_OverSample = True
Generate_NoLabel = False
Add_Pseudo_Label = True

upper_limit = 100
lower_limit = 1
random_limit = 0.5  # 随机拓展 0.5:15W->21W

if Add_Pseudo_Label:
    print('add pseudo label to train 2020 txt')
    data_dir = '../../data/train'
    os.system('cp label_source.txt label.txt')
    sourece_label_file = '/label.txt'
    pseudo_label_file = '/pse_label.txt'
    sourece_label = open(data_dir+sourece_label_file,'a')
    pseudo_label = open(data_dir+pseudo_label_file,'r')

    a = pseudo_label.read()
    sourece_label.write(a)

    sourece_label.close()
    pseudo_label.close()

if Train2020_OverSample:
    print('over sample for train 2020......')
    data_dir = '../../data/train'
    filename = data_dir + '/label.txt'
    count_image = defaultdict(list)  # 所有的数据集 dict {'0':['1.png','2.png',...],'1':[],...}
    with open(filename, 'r') as file_to_read:
        while True:
            lines = file_to_read.readline()
            if not lines:
                break
            (img_name, img_label) = lines.strip().split(':')
            count_image[img_label].append(img_name)
        pid_num = {i: len(count_image[i]) for i in count_image.keys()}
        for i in pid_num:
            if (pid_num[i]<upper_limit) and (pid_num[i]>lower_limit) and (random.random()<random_limit):
                count_image[i].extend(count_image[i])
        file = open(data_dir + '/label.txt', 'w')
        for i in count_image.keys():
            for filename in count_image[i]:
                file.write(str(filename+':'+i + '\n'))
        file.close()

if Generate_NoLabel:
    print("Generate no label TXT for train 2020")
    "生成无标签数据集——标签默认给1"
    data_dir = '../../data/train'
    filename = data_dir + '/label_source.txt'
    all_image_list = os.listdir(data_dir+'/images')
    pseudo_label = []
    label_image = []
    with open(filename, 'r') as file_to_read:
        while True:
            lines = file_to_read.readline()
            if not lines:
                break
            (img_name, img_label) = lines.strip().split(':')
            label_image.append(img_name)
    for all_image in tqdm(all_image_list):
        if all_image not in label_image:
            pseudo_label.append(all_image)

    file = open(data_dir +'/pseudo_label.txt', 'w')
    for i in pseudo_label:
        file.write(str(i+':'+ '1' + '\n'))
    file.close()


if Train2019_OverSample:
    print('over sample for train 2019......')
    data_dir = '../../data/2019/round2'
    filename = data_dir + '/train_list_source.txt'
    count_image = defaultdict(list)  # 所有的数据集 dict {'0':['1.png','2.png',...],'1':[],...}
    with open(filename, 'r') as file_to_read:
        while True:
            lines = file_to_read.readline()
            if not lines:
                break
            (img_name, img_label) = lines.strip().split(' ')
            count_image[img_label].append(img_name)
        pid_num = {i: len(count_image[i]) for i in count_image.keys()}
        for i in pid_num:
            if (pid_num[i]<upper_limit) and (pid_num[i]>lower_limit) and (random.random()<random_limit):
                count_image[i].extend(count_image[i])
        file = open(data_dir +'/train_list.txt', 'w')
        for i in count_image.keys():
            for filename in count_image[i]:
                file.write(str(filename+' '+i + '\n'))
        file.close()